#
from typing import Dict
import json
from apps.glm4.glm4_engine import Glm4Engine

class Glm4App(object):
    def __init__(self):
        self.name = 'apps.glm4.glm4_app.Gl4App'

    @staticmethod
    def startup(params:Dict = {}) -> None:
        print(f'GLM4 App v0.0.2')
        if params['task'] == 1:
            query = '请介绍基于GLM-4的多Agent开发技术<|endoftext|>'
            Glm4Engine.infer(query=query, model=Glm4Engine.BASE_MODEL)
        elif params['task'] == 2:
            query = '请介绍智能化雷达核心技术'
            Glm4Engine.infer(query=query, model=Glm4Engine.CHAT_MODEL)
        elif params['task'] == 3:
            query = '类型#裙*版型#宽松*材质#雪纺*风格#清新*裙型#a字*裙长#连衣裙'
            print(f'Reference: 踩着轻盈的步伐享受在午后的和煦风中，让放松与惬意感为你免去一身的压力与束缚，仿佛要将灵魂也寄托在随风摇曳的雪纺连衣裙上，吐露出<UNK>微妙而又浪漫的清新之意。宽松的a字版型除了能够带来足够的空间，也能以上窄下宽的方式强化立体层次，携带出自然优雅的曼妙体验。')
            Glm4Engine.infer(query=query, model=Glm4Engine.B0414_MODEL)
        elif params['task'] == 4: # Prepare the dataset
            Glm4App.prepare_AdvertiseGen()
        elif params['task'] == 5: # Infer with 9b 0414
            query = '类型#裙*版型#宽松*材质#雪纺*风格#清新*裙型#a字*裙长#连衣裙'
            print(f'Reference: 踩着轻盈的步伐享受在午后的和煦风中，让放松与惬意感为你免去一身的压力与束缚，仿佛要将灵魂也寄托在随风摇曳的雪纺连衣裙上，吐露出<UNK>微妙而又浪漫的清新之意。宽松的a字版型除了能够带来足够的空间，也能以上窄下宽的方式强化立体层次，携带出自然优雅的曼妙体验。')
            Glm4Engine.infer(query=query, model=Glm4Engine.B0414_MODEL)
        elif params['task'] == 6: # Infer with LoRA of 9b 0414
            print(f'LoRA finetune infer')
            query = '类型#裙*版型#宽松*材质#雪纺*风格#清新*裙型#a字*裙长#连衣裙'
            print(f'Reference: 踩着轻盈的步伐享受在午后的和煦风中，让放松与惬意感为你免去一身的压力与束缚，仿佛要将灵魂也寄托在随风摇曳的雪纺连衣裙上，吐露出<UNK>微妙而又浪漫的清新之意。宽松的a字版型除了能够带来足够的空间，也能以上窄下宽的方式强化立体层次，携带出自然优雅的曼妙体验。')
            Glm4Engine.infer(query=query, model=Glm4Engine.B0414_LORA_MODEL)


    @staticmethod
    def prepare_AdvertiseGen() -> None:
        # 转训练集
        src_fn = 'llms/datasets/AdvertiseGen/train.json'
        dst_fn = 'llms/datasets/glm4_AdvertiseGen/train.jsonl'
        Glm4App._prepare_AdvertiseGen(src_fn=src_fn, dst_fn=dst_fn)
        # 转测试集
        src_fn = 'llms/datasets/AdvertiseGen/dev.json'
        dst_fn = 'llms/datasets/glm4_AdvertiseGen/dev.jsonl'
        Glm4App._prepare_AdvertiseGen(src_fn=src_fn, dst_fn=dst_fn)

    @staticmethod
    def _prepare_AdvertiseGen(src_fn:str, dst_fn:str) -> None:
        recs = []
        num = 0
        with open(src_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                rec = json.loads(row)
                recs.append(rec)
                num += 1
                if num % 100 == 0:
                    print(f'读入{num}条')
        total = num
        num = 0
        with open(dst_fn, 'w', encoding='utf-8') as wfd:
            for rec in recs:
                item = {
                    'messages': [
                        {
                            'role': 'user',
                            'content': rec['content']
                        },
                        {
                            'role': 'assistant',
                            'content': rec['summary']
                        }
                    ]
                }
                wfd.write(f'{json.dumps(item, ensure_ascii=False)}\n')
                num += 1
                if num % 100 == 0:
                    print(f'处理完成{num}/{total}条')